import torch as th
import torch.nn as nn
import torch.nn.functional as F

from ..maxvit.maxvit import MLP, get_act_layer


class Conv_BN_ReLU(nn.Module):
    def __init__(self, input_channels, output_channles,
                 kernel_size=3, stride=1, padding=1,
                 bias=True):
        super().__init__()
        self.CBR = nn.Sequential(
            nn.Conv2d(input_channels, output_channles,
                      kernel_size=kernel_size, stride=stride, padding=padding,
                      bias=bias),
            nn.BatchNorm2d(output_channles),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.CBR(x)


class Cat_Fusion(nn.Module):
    def __init__(self, img_feature_channels, ev_feature_channels,
                 output_channels):
        super().__init__()

        assert img_feature_channels == ev_feature_channels
        self.cat_CBR = Conv_BN_ReLU(output_channels * 2, output_channels,
                                    kernel_size=3, stride=1, padding=1)

    def forward(self, frame_coarse_feature, event_coarse_feature):
        output_feature = th.cat([frame_coarse_feature, event_coarse_feature], dim=1)
        output_feature = self.cat_CBR(output_feature)

        return output_feature


class Feature_fusion(nn.Module):
    def __init__(self, img_feature_channels, ev_feature_channels,
                 output_channels):
        super().__init__()

        self.frame_c = output_channels
        self.event_c = output_channels

        self.frame_CBR = Conv_BN_ReLU(img_feature_channels, self.frame_c,
                                      kernel_size=3, stride=1, padding=1)
        self.event_CBR = Conv_BN_ReLU(ev_feature_channels, self.event_c,
                                      kernel_size=3, stride=1, padding=1)
        self.cat_CBR = Conv_BN_ReLU(output_channels * 2, output_channels,
                                    kernel_size=3, stride=1, padding=1)

    def forward(self, frame_coarse_feature, event_coarse_feature):
        frame_feature = self.frame_CBR(frame_coarse_feature)  # [B, 32, H, W]
        event_feature = self.event_CBR(event_coarse_feature)  # [B, 72, H, W]

        frame_feature_weighted = frame_feature + frame_feature * event_feature
        event_feature_weighted = event_feature + event_feature * frame_feature

        output_feature = th.cat([frame_feature_weighted, event_feature_weighted], dim=1)
        output_feature = self.cat_CBR(output_feature)

        return output_feature


class Seletive_Feature_fusion(nn.Module):
    def __init__(self, img_feature_channels, ev_feature_channels,
                 output_channels):
        super().__init__()

        self.frame_c = output_channels
        self.event_c = output_channels

        self.frame_CBR = Conv_BN_ReLU(img_feature_channels, self.frame_c,
                                      kernel_size=3, stride=1, padding=1)
        self.event_CBR = Conv_BN_ReLU(ev_feature_channels, self.event_c,
                                      kernel_size=3, stride=1, padding=1)

        self.weight_frame_CBR = Conv_BN_ReLU(output_channels * 2, output_channels,
                                             kernel_size=3, stride=1, padding=1)

        self.weight_event_CBR = Conv_BN_ReLU(output_channels * 2, output_channels,
                                             kernel_size=3, stride=1, padding=1)

        # self.output_MLP = MLP(
        #             dim = output_channels,
        #             channel_last=True,
        #             expansion_ratio = 4,
        #             act_layer = get_act_layer('gelu'),
        #             gated = 0,
        #             bias = True,
        #             drop_prob = 0)

        self.output_CBR = Conv_BN_ReLU(output_channels, output_channels,
                                       kernel_size=3, stride=1, padding=1)

    def forward(self, frame_coarse_feature, event_coarse_feature):
        # enhance the similar feature

        frame_feature = self.frame_CBR(frame_coarse_feature)
        event_feature = self.event_CBR(event_coarse_feature)

        frame_feature_weighted = frame_feature + frame_feature * event_feature
        event_feature_weighted = event_feature + event_feature * frame_feature

        cat_feature = th.cat([frame_feature_weighted, event_feature_weighted], dim=1)

        # data dependently select feature

        weight_frame = self.weight_frame_CBR(cat_feature)
        weight_event = self.weight_event_CBR(cat_feature)

        output_feature = weight_frame * frame_feature_weighted + weight_event * event_feature_weighted
        output_feature = self.output_CBR(output_feature)

        return output_feature

